from flask import Blueprint, request, jsonify, session, current_app
from werkzeug.utils import secure_filename
from models.user import User
from models.image_task import ImageTask
from models.style_template import StyleTemplate
from models.ai_config import AIConfig
from tasks.image_generation import generate_image_task
from app import db
import uuid
import os
from PIL import Image

image_bp = Blueprint('image', __name__)

ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg', 'gif', 'webp'}

def allowed_file(filename):
    return '.' in filename and \
           filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS

def resize_image(image_path, max_size=1024):
    """调整图片大小"""
    with Image.open(image_path) as img:
        # 保持宽高比，限制最大尺寸
        img.thumbnail((max_size, max_size), Image.Resampling.LANCZOS)
        
        # 转换为RGB（如果是RGBA）
        if img.mode in ('RGBA', 'P'):
            img = img.convert('RGB')
        
        # 保存调整后的图片
        img.save(image_path, 'JPEG', quality=85, optimize=True)

@image_bp.route('/upload', methods=['POST'])
def upload_image():
    """上传图片"""
    user_id = session.get('user_id')
    if not user_id:
        return jsonify({'error': '请先登录'}), 401
    
    user = User.query.get(user_id)
    if not user:
        return jsonify({'error': '用户不存在'}), 404
    
    # 检查用户是否还能生成图片
    if not user.can_generate():
        return jsonify({'error': '今日生成次数已用完，请明天再试'}), 403
    
    if 'file' not in request.files:
        return jsonify({'error': '没有选择文件'}), 400
    
    file = request.files['file']
    if file.filename == '':
        return jsonify({'error': '没有选择文件'}), 400
    
    if file and allowed_file(file.filename):
        try:
            # 生成安全的文件名
            filename = secure_filename(file.filename)
            file_ext = filename.rsplit('.', 1)[1].lower()
            unique_filename = f"{uuid.uuid4().hex}.{file_ext}"
            
            # 保存文件
            upload_folder = current_app.config['UPLOAD_FOLDER']
            file_path = os.path.join(upload_folder, unique_filename)
            file.save(file_path)
            
            # 调整图片大小
            resize_image(file_path, max_size=1024)
            
            # 生成访问URL
            file_url = f"/uploads/{unique_filename}"
            
            return jsonify({
                'message': '文件上传成功',
                'file_url': file_url,
                'filename': unique_filename
            })
            
        except Exception as e:
            return jsonify({'error': f'文件上传失败: {str(e)}'}), 500
    
    return jsonify({'error': '不支持的文件类型'}), 400

@image_bp.route('/generate', methods=['POST'])
def generate_image():
    """生成图片"""
    user_id = session.get('user_id')
    if not user_id:
        return jsonify({'error': '请先登录'}), 401
    
    user = User.query.get(user_id)
    if not user:
        return jsonify({'error': '用户不存在'}), 404
    
    # 检查用户是否还能生成图片
    if not user.can_generate():
        return jsonify({'error': '今日生成次数已用完，请明天再试'}), 403
    
    try:
        data = request.get_json()
        original_image_url = data.get('original_image_url')
        style_template_id = data.get('style_template_id')
        custom_prompt = data.get('custom_prompt', '')
        
        if not original_image_url or not style_template_id:
            return jsonify({'error': '缺少必要参数'}), 400
        
        # 验证风格模板
        style_template = StyleTemplate.query.get(style_template_id)
        if not style_template or not style_template.is_active:
            return jsonify({'error': '风格模板不存在或已禁用'}), 404
        
        # 检查VIP权限
        if style_template.is_premium and user.vip_level == 0:
            return jsonify({'error': '该风格需要VIP权限'}), 403
        
        # 获取可用的AI配置
        ai_config = AIConfig.query.filter_by(is_active=True).order_by(AIConfig.priority.desc()).first()
        if not ai_config:
            return jsonify({'error': '暂无可用的AI服务，请稍后再试'}), 503
        
        # 创建生成任务
        task_id = str(uuid.uuid4())
        
        # 构建提示词
        positive_prompt = style_template.positive_prompt
        if custom_prompt:
            positive_prompt = f"{positive_prompt}, {custom_prompt}"
        
        image_task = ImageTask(
            id=task_id,
            user_id=user_id,
            style_template_id=style_template_id,
            ai_config_id=ai_config.id,
            original_image_url=original_image_url,
            positive_prompt=positive_prompt,
            negative_prompt=style_template.negative_prompt,
            custom_prompt=custom_prompt,
            ai_model=ai_config.model_name,
            width=style_template.recommended_width,
            height=style_template.recommended_height,
            steps=style_template.recommended_steps,
            cfg_scale=style_template.recommended_cfg_scale
        )
        
        db.session.add(image_task)
        db.session.commit()
        
        # 启动异步任务
        generate_image_task.delay(task_id)
        
        # 增加使用次数
        user.used_quota += 1
        style_template.increment_usage()
        ai_config.increment_usage()
        
        db.session.commit()
        
        return jsonify({
            'message': '任务创建成功，开始生成图片',
            'task_id': task_id,
            'task': image_task.to_dict()
        })
        
    except Exception as e:
        db.session.rollback()
        return jsonify({'error': f'生成失败: {str(e)}'}), 500

@image_bp.route('/task/<task_id>', methods=['GET'])
def get_task_status(task_id):
    """获取任务状态"""
    user_id = session.get('user_id')
    if not user_id:
        return jsonify({'error': '请先登录'}), 401
    
    task = ImageTask.query.filter_by(id=task_id, user_id=user_id).first()
    if not task:
        return jsonify({'error': '任务不存在'}), 404
    
    return jsonify({
        'task': task.to_dict()
    })

@image_bp.route('/history', methods=['GET'])
def get_history():
    """获取用户的生成历史"""
    user_id = session.get('user_id')
    if not user_id:
        return jsonify({'error': '请先登录'}), 401
    
    page = request.args.get('page', 1, type=int)
    per_page = request.args.get('per_page', 10, type=int)
    
    pagination = ImageTask.query.filter_by(user_id=user_id)\
        .order_by(ImageTask.created_at.desc())\
        .paginate(page=page, per_page=per_page, error_out=False)
    
    tasks = [task.to_dict() for task in pagination.items]
    
    return jsonify({
        'tasks': tasks,
        'total': pagination.total,
        'pages': pagination.pages,
        'current_page': page,
        'has_next': pagination.has_next,
        'has_prev': pagination.has_prev
    })